#!/usr/bin/env python

"""
Script to find all the genuses present in the ARB database, get all of their
HMMs, and save all of those HMMs to a single file. Once all those HMMs are in a
single file, they're going to be pressed by hmmbuild.

"""

from ssummo import CONFIG
from ssummo.traverse import dict_walk

import os
import subprocess
import sys
from Queue import Queue
from threading import Thread
try:
    import cPickle as pickle
except ImportError:
    import pickle

import MySQLdb

class FindHmmPath(Thread):
    """This thread continually loops through the dictionary
    looking for path names that match the OTU names piped over.
    When the OTU name & parent name match, the path to that
    HMM is passed back to the main thread.
    """
    def __init__(self, tdict, inQ, outQ):
        Thread.__init__(self)
        self.tdict = tdict
        self.inQ = inQ
        self.outQ = outQ
#        self.l = lock
#        self.sem = sem

    def run(self):
        obj = self.inQ.get( )
        self.name, self.parentName = obj
        while True:
            for path, node in dict_walk(CONFIG.arbDBdir, self.tdict):
                taxPath = path.rsplit(os.path.sep, 2)
                if taxPath[-1] == self.name and taxPath[-2] == self.parentName:
                    self.inQ.task_done()
                    self.outQ.put(os.path.join(path, taxPath[-1] + '.hmm'))
                    #^^^ Pass path to HMM back to main process.
                    obj = self.inQ.get( )
                    if obj == 'END':
                        self.inQ.task_done()
                        break
                    else:
                        self.name, self.parentName = obj
            if obj == 'END':
                break

def getGenuses(domain):
    cnx = MySQLdb.connect(user=CONFIG.user, host=CONFIG.host,
                          db=CONFIG.db, passwd=CONFIG.passwd)
    cur = cnx.cursor()
    cmd = 'SELECT tax_ID,Name,parentName FROM {0} WHERE rank="genus" OR ' + \
          'rank="Genus";'.format(domain)
    cur.execute(cmd)
    try:
        for row in cur.fetchall():
            yield row
    except TypeError:
        print(row)
        raise

def pressHMM(inName):
    return subprocess.Popen(['hmmpress', inName], shell=False,
                            stdout=subprocess.PIPE, stderr=subprocess.PIPE )

def main():
    if len(sys.argv) == 1:
        print('Going to press all genera from each domain, but...')
        msg = 'Which domain?? [E]ukaryotes, [P]rokaryotes? or [c]ancel. [p]'
        domain = raw_input(msg)
        if len(domain) == 0 or domain.startswith('p') or domain.startswith('P'):
            domain = 'Prokaryotes'
        elif domain.startswith('e') or domain.startswith('E'):
            domain = 'Eukaryotes'
        else:
            print("Got {0}.\nExiting".format(domain))
            exit()
    else:
        domain = sys.argv[1]
    outName = CONFIG.taxIndex[:CONFIG.taxIndex.rfind('.')] + '_to_press.hmm'
    #^^^ Where to save all the genus HMMs.
    if domain != 'Eukaryotes' and domain != 'Prokaryotes':
        sys.stderr.write("\n\n")
        sys.stderr.write("Don't recognise this as a domain:-\n  {0}\n"
                         .format(domain))
        sys.stderr.write("Exiting..\n" )
        exit()
    with file(CONFIG.taxIndex, 'rb') as taxFile:
        sys.stdout.write( 'Loading the complete taxonomy index. . . ')
        tdict = pickle.load(taxFile)  # Load the taxonomy index.
        print('Done.')
    outFile = file(outName, 'w')    # Open output file.
    inQ = Queue()
    outQ = Queue()
    finder = FindHmmPath(tdict, inQ, outQ)   # Initiate process.
    finder.start()                           # run it.
    for taxID, OTU, parent in getGenuses( domain ):
        inQ.put( (OTU, parent) )
        path = outQ.get( )
        # Open [genus name].hmm
        with file(path, 'r') as infile:
            # Append the hmm file's contents to the output file [outName].
            outFile.write(infile.read())
        outQ.task_done()
    inQ.put('END')  # Close pipes.
    outQ.join()
    inQ.join()
    finder.join()
    outFile.close()
    print('Found all {0} at genus. Pressing HMMs now into:-\n{1}.*.'
          .format(domain, outName))
    print('Pressing all genera within {0} into:-\n{1}'
          .format(domain, outName))  # Tell user what's being made.
    process = pressHMM(outName)      # Press the HMM.
    return_code = process.wait()
    if return_code != 0:
        sys.stderr.write( 'hmmpress failed with this error:-\n' )
        sys.stderr.write( process.stdout.read() + process.stderr.read() )
    os.remove( outName )

if __name__ == '__main__':
    main()
